import tensorflow as tf
import numpy as np

# 加载MNIST数据集
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# 对数据进行预处理，将像素值大于阈值的设为255，小于阈值的设为0，并进行归一化
threshold = 128
x_train = np.where(x_train > threshold, 255, 0)
x_test = np.where(x_test > threshold, 255, 0)
x_train, x_test = x_train / 255.0, x_test / 255.0

print(np.ndim(x_train))
# 构建模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),  # 将28x28的图像展平成一维向量
    tf.keras.layers.Dense(128, activation='relu'),   # 全连接层，128个神经元，ReLU激活函数
    tf.keras.layers.Dropout(0.2),                   # Dropout层，防止过拟合
    tf.keras.layers.Dense(10),                      # 输出层，有10个神经元，对应10个类别
    tf.keras.layers.Softmax()                       # 在最后一层添加softmax激活函数，输出概率分布
])

# 打印模型的预测输出和softmax概率
predictions = model(x_train[:1]).numpy()
print("模型预测输出：")
print(predictions)
print("softmax概率输出：")
print(tf.nn.softmax(predictions).numpy())

# 定义损失函数
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 计算模型在一个样本上的损失值
loss = loss_fn(y_train[:1], predictions).numpy()
print("模型在一个样本上的损失值：", loss)

# 编译模型，指定优化器、损失函数和评估指标
model.compile(optimizer='RMSprop',
              loss=loss_fn,
              metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=10)

# 保存训练好的模型
model.save("path_to_save_model.h5")

# 在测试集上进行评估
print("在测试集上评估模型：")
eval_loss, eval_accuracy = model.evaluate(x_test, y_test, verbose=2)
print("测试集上的损失值：", eval_loss)
print("测试集上的准确率：", eval_accuracy)

# 进行单个样本的预测
sample_index = 10
single_image = x_test[sample_index]
single_image = np.expand_dims(single_image, axis=0)  # 添加一个维度以适应模型输入
predictions = model.predict(single_image)

# 获取预测结果
predicted_label = np.argmax(predictions[0])  # 获取概率最大的类别索引

print("预测结果为：", predicted_label)
print("正确结果为：", y_test[sample_index])
